from langchain.prompts import SemanticSimilarityExampleSelector
from langchain_community.vectorstores import Chroma
from langchain_openai import OpenAIEmbeddings
from langchain_openai import ChatOpenAI
from langchain.prompts import FewShotChatMessagePromptTemplate
from langchain.prompts import ChatPromptTemplate


api_key = "sk-6S0PtpNia71gjcfwSsDPsJ9mGqsVPr2XRQzAx1dHbJS7RW4t"
api_base="https://chatapi.littlewheat.com/v1"

examples = [
    # 数学推理
    {
        "question": "小明的妈妈给了他10块钱去买文具，如果一支笔3块钱，小明最多能买几支笔？",
        "answer": "小明有10块钱，每支笔3块钱，所以他最多能买3支笔，因为3*3=9，剩下1块钱不够再买一支笔。因此答案是3支。"
    },
    {
        "question": "一个篮球队有12名球员，如果教练想分成两个小组进行训练，每组需要有多少人？",
        "answer": "篮球队总共有12名球员，分成两个小组，每组有12/2=6名球员。因此每组需要有6人。"
    },
    # 逻辑推理
    {
        "question": "如果所有的猫都怕水，而Tom是一只猫，请问Tom怕水吗？",
        "answer": "根据题意，所有的猫都怕水，因此作为一只猫的Tom也会怕水。所以答案是肯定的，Tom怕水。"
    },
    {
        "question": "在夏天，如果白天温度高于30度，夜晚就会很凉爽。今天白天温度是32度，请问今晚会凉爽吗？",
        "answer": "根据题意，只要白天温度高于30度，夜晚就会很凉爽。今天白天的温度是32度，超过了30度，因此今晚会凉爽。"
    },
    # 常识问题
    {
        "question": "地球绕太阳转一圈需要多久？",
        "answer": "地球绕太阳转一圈大约需要365天，也就是一年的时间。"
    },
    {
        "question": "水的沸点是多少摄氏度？",
        "answer": "水的沸点是100摄氏度。"
    },
    # 文化常识
    {
        "question": "中国的首都是哪里？",
        "answer": "中国的首都是北京。"
    },
    {
        "question": "世界上最长的河流是哪一条？",
        "answer": "世界上最长的河流是尼罗河。"
    },
]

to_vectorizer = [" ".join(example.values()) for example in examples]
#print(to_vectorizer)

embeddings = OpenAIEmbeddings(model="text-embedding-ada-002",api_key=api_key ,base_url=api_base)
vectorstore = Chroma.from_texts(to_vectorizer,embeddings,metadatas=examples)

example_selector = SemanticSimilarityExampleSelector(
    vectorstore=vectorstore,
    k=2
)
selector=example_selector.select_examples({"input":"广东的省份是哪个城市?"})
print(selector)

few_shot_chat_prompt = FewShotChatMessagePromptTemplate(
    input_variables=["input"],
    example_selector=example_selector,
    example_prompt=ChatPromptTemplate.from_messages(
        [
            ("human","{question}"),
            ("ai","{answer}")
        ]
    )
)

chat_prompt = few_shot_chat_prompt.format(input="月亮每天什么时候出现")
print(chat_prompt)

final_prompt = ChatPromptTemplate.from_messages(
    [
        ("system","你是一个无所不能的人，无论什么问题都可以回答。"),
        few_shot_chat_prompt,
        ("human","{input}")
    ]
)
print(final_prompt)

chat = ChatOpenAI(model="gpt-3.5-turbo",api_key=api_key ,base_url=api_base)
chain = final_prompt | chat
response=chain.invoke({"input":"月亮每天什么时候出现"})
print(response.content)